'''
VOC xml to SDL
'''
import os
from PIL import Image
from utils import parse_voc_xml
class_names = ['papercup']


# resize img
os.environ["data_root"] = data_root = "../data/"
os.environ["im2rec"] = "python /home/ubuntu/anaconda3/lib/python3.6/site-packages/mxnet/tools/im2rec.py"
resize = (512, 512)
os.environ["resize"] = resize_str = str(resize[0]) + "_" + str(resize[1])
for imgname in os.listdir(data_root + '/img'):
        imgpath = data_root + "/img/" + imgname
        img = Image.open(imgpath)
        img = img.resize(resize, Image.BILINEAR)
        img.save(data_root + 'img%d_%d' % (resize[0], resize[1]) + "/" + imgname)

# generate lst
os.system('$im2rec --list True --train-ratio 0.9 ${data_root}/rec/img_$resize ${data_root}/img%d_%d' % (resize[0], resize[1]))


# modify lst
'''
lst格式：
idx \t header_width \t label_width \t [labels] \t filename
label_width为每个label的宽度。header_width为label之前idx之后的数据宽度，一般为2。
label: class_idx \t xmin \t ymin \t xmax \t ymax
'''
new_lst_content = ''
with open(data_root + '/rec/img_%s_train.lst' % (resize_str)) as f:
    contents = f.read().split('\n')
    for content in contents:
        if content == '':
            break
        content = content.split('\t')
        idx = content[0]
        file = content[-1][:-3] + 'xml'

        bndboxs, names, filename = parse_voc_xml(data_root + '/labels/' + file)

        data = idx + '\t2\t5\t'
        for bndbox, name in zip(bndboxs, names):
            data +=  "%d\t%f\t%f\t%f\t%f\t" % (class_names.index(name), bndbox[0], bndbox[1], bndbox[2], bndbox[3])
        data += filename
        print(data)
        new_lst_content += data + '\n'
with open(data_root + '/rec/img_%s_train.lst' % (resize_str), 'w') as f:
    f.write(new_lst_content)



new_lst_content = ''
with open(data_root + '/rec/img_%s_val.lst' % (resize_str)) as f:
    contents = f.read().split('\n')
    for content in contents:
        if content == '':
            break
        content = content.split('\t')
        idx = content[0]
        file = content[-1][:-3] + 'xml'

        bndboxs, names, filename = parse_voc_xml(data_root + '/labels/' + file)

        data = idx + '\t2\t5\t'
        for bndbox, name in zip(bndboxs, names):
            data +=  "%d\t%f\t%f\t%f\t%f\t" % (class_names.index(name), bndbox[0], bndbox[1], bndbox[2], bndbox[3])
        data += filename
        print(data)
        new_lst_content += data + '\n'
with open(data_root + '/rec/img_%s_val.lst' % (resize_str), 'w') as f:
    f.write(new_lst_content)

# generate rec file
os.system('$im2rec --num-thread 10 --pass-through 1 --pack-label True $data_root/rec/img_$resize $data_root/img$resize --encoding=.jpg --quality 100')

# 计算mean
# mean = cal_mean('../data/img512_512')
# print(mean)
# with open('../data/mean_512_512.txt', 'w') as f:
#     f.write("%f %f %f" % (mean[0], mean[1], mean[2]))